from typing import List, Tuple
from langchain_core.documents import Document
from sentence_transformers import CrossEncoder


class ReRanker:
    def __init__(self, model_path: str):

        try:
            self.model = CrossEncoder(model_path)
        except Exception as e:
            self.model = None
            print("[WARNING] reranker load error.")

    def rerank(self, query: str, passages: List[Document], top_k: int = 5) -> List[Tuple[Document, float]]:
        """
        输入query和召回内容列表(注意：列表中的内容应该是扩充上下文之后的内容，或者完整的qa)
        :param query:
        :param passages: 召回的文本列表，一般召回20-30条内容就能比较全了
        :param top_k: 最终喂给LLM的引用条数，测试发现3-5条的效果就够用了，超过10条提升不大
        :return: 文档列表，其中每个文档包含原始内容和reranker分数，这个分数越高，代表越相似
        """
        ranked_li = []
        for doc in passages:
            score = self.model.predict([query, doc.page_content])
            ranked_li.append([doc, score])
        return [(doc, _) for doc, _ in sorted(ranked_li, key=lambda x: x[1], reverse=True)[:top_k]]

